{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 27,
   "id": "f2c3496f-5fc3-4975-a812-f5a066dda84f",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import cv2\n",
    "import json\n",
    "import glob\n",
    "import shutil\n",
    "import random\n",
    "import numpy as np\n",
    "import os.path as osp\n",
    "import fiftyone as fo\n",
    "\n",
    "from tqdm import tqdm\n",
    "from pycocotools.coco import COCO\n",
    "from pycocotools import mask as maskUtils\n",
    "\n",
    "from my_class import *"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "id": "1769f529-4f05-43b8-9ddc-37edc4367ff0",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "nums of batch:  1\n"
     ]
    }
   ],
   "source": [
    "metric_num = \"test0425exp1\"\n",
    "root_dir = f\"/home/ray/input/AI_Image/202504/{metric_num}/\"\n",
    "# root_dir = \"/home/user/workspace/datasets/tmp/res\"\n",
    "batch_nums = os.listdir(root_dir)\n",
    "\n",
    "# 随机10个批次\n",
    "# batch_nums = random.sample(batch_nums, 2)\n",
    "print(\"nums of batch: \", len(batch_nums))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "id": "89a44a22-816c-48df-8314-7b9a8e3cfff0",
   "metadata": {},
   "outputs": [],
   "source": [
    "dataset_name=\"nike_infer\"\n",
    "fo.delete_datasets('*')\n",
    "\n",
    "dataset = fo.Dataset(name=dataset_name)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 38,
   "id": "f9cc56cd-792a-4e18-906c-3a1a7fbe3064",
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 20/20 [00:00<00:00, 527.24it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "/home/ray/input/AI_Image/202504/test0425exp1/20250425-1/1/Left/test0425exp1_20250425-1_1_Left_Camra01_blue.png\n",
      "1\n",
      "/home/ray/input/AI_Image/202504/test0425exp1/20250425-1/1/Left/test0425exp1_20250425-1_1_Left_Camra01_white.json\n",
      "/home/ray/input/AI_Image/202504/test0425exp1/20250425-1/1/Left/test0425exp1_20250425-1_1_Left_Camra02_blue.png\n",
      "2\n",
      "/home/ray/input/AI_Image/202504/test0425exp1/20250425-1/1/Left/test0425exp1_20250425-1_1_Left_Camra03_blue.png\n",
      "2\n",
      "/home/ray/input/AI_Image/202504/test0425exp1/20250425-1/1/Left/test0425exp1_20250425-1_1_Left_Camra04_blue.png\n",
      "2\n",
      "/home/ray/input/AI_Image/202504/test0425exp1/20250425-1/1/Left/test0425exp1_20250425-1_1_Left_Camra05_blue.png\n",
      "2\n",
      "/home/ray/input/AI_Image/202504/test0425exp1/20250425-1/1/Right/test0425exp1_20250425-1_1_Right_Camra01_blue.png\n",
      "1\n",
      "/home/ray/input/AI_Image/202504/test0425exp1/20250425-1/1/Right/test0425exp1_20250425-1_1_Right_Camra01_white.json\n",
      "/home/ray/input/AI_Image/202504/test0425exp1/20250425-1/1/Right/test0425exp1_20250425-1_1_Right_Camra02_blue.png\n",
      "2\n",
      "/home/ray/input/AI_Image/202504/test0425exp1/20250425-1/1/Right/test0425exp1_20250425-1_1_Right_Camra03_blue.png\n",
      "2\n",
      "/home/ray/input/AI_Image/202504/test0425exp1/20250425-1/1/Right/test0425exp1_20250425-1_1_Right_Camra04_blue.png\n",
      "2\n",
      "/home/ray/input/AI_Image/202504/test0425exp1/20250425-1/1/Right/test0425exp1_20250425-1_1_Right_Camra05_blue.png\n",
      "2\n",
      "/home/ray/input/AI_Image/202504/test0425exp1/20250425-1/2/Left/test0425exp1_20250425-1_2_Left_Camra01_blue.png\n",
      "1\n",
      "/home/ray/input/AI_Image/202504/test0425exp1/20250425-1/2/Left/test0425exp1_20250425-1_2_Left_Camra01_white.json\n",
      "/home/ray/input/AI_Image/202504/test0425exp1/20250425-1/2/Left/test0425exp1_20250425-1_2_Left_Camra02_blue.png\n",
      "2\n",
      "/home/ray/input/AI_Image/202504/test0425exp1/20250425-1/2/Left/test0425exp1_20250425-1_2_Left_Camra03_blue.png\n",
      "2\n",
      "/home/ray/input/AI_Image/202504/test0425exp1/20250425-1/2/Left/test0425exp1_20250425-1_2_Left_Camra04_blue.png\n",
      "2\n",
      "/home/ray/input/AI_Image/202504/test0425exp1/20250425-1/2/Left/test0425exp1_20250425-1_2_Left_Camra05_blue.png\n",
      "2\n",
      "/home/ray/input/AI_Image/202504/test0425exp1/20250425-1/2/Right/test0425exp1_20250425-1_2_Right_Camra01_blue.png\n",
      "1\n",
      "/home/ray/input/AI_Image/202504/test0425exp1/20250425-1/2/Right/test0425exp1_20250425-1_2_Right_Camra01_white.json\n",
      "/home/ray/input/AI_Image/202504/test0425exp1/20250425-1/2/Right/test0425exp1_20250425-1_2_Right_Camra02_blue.png\n",
      "2\n",
      "/home/ray/input/AI_Image/202504/test0425exp1/20250425-1/2/Right/test0425exp1_20250425-1_2_Right_Camra03_blue.png\n",
      "2\n",
      "/home/ray/input/AI_Image/202504/test0425exp1/20250425-1/2/Right/test0425exp1_20250425-1_2_Right_Camra04_blue.png\n",
      "2\n",
      "/home/ray/input/AI_Image/202504/test0425exp1/20250425-1/2/Right/test0425exp1_20250425-1_2_Right_Camra05_blue.png\n",
      "2\n",
      "20\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "def parse_result(annFile):\n",
    "    try:\n",
    "        with open(annFile, \"r\") as f:\n",
    "            data = json.load(f)\n",
    "        print(len(data))\n",
    "        if len(data) == 1:\n",
    "            print(annFile)\n",
    "            data = data[0]\n",
    "        else:\n",
    "            data = data[1]\n",
    "        predictions = data[\"predictions\"]\n",
    "        fg = data['offset'][:4]\n",
    "        result = [Anomaly(data['offset'][4], fg, None, \"shoe\")]\n",
    "    \n",
    "        for item in predictions:\n",
    "            flaw = Anomaly(item[\"scores\"], item[\"bboxes\"], item[\"segmentation\"], item[\"category_name\"])\n",
    "            result.append(flaw)\n",
    "        return result\n",
    "    except Exception as e:\n",
    "        print(e)\n",
    "        return None\n",
    "\n",
    "samples = []\n",
    "\n",
    "light = \"white\"\n",
    "for batch in batch_nums:\n",
    "    img_file = glob.glob(f\"{osp.join(root_dir, batch)}/**/*{light}.png\", recursive=True)\n",
    "    img_shape = (4096, 3000)\n",
    "    img_w = img_shape[0]\n",
    "    img_h = img_shape[1]\n",
    "    \n",
    "    for img in tqdm(img_file):\n",
    "        print(img)\n",
    "        pred_file = img.replace(f\"{light}.png\", \"white.json\")\n",
    "        flaws = parse_result(pred_file)\n",
    "        detections = []\n",
    "        segmentations = []\n",
    "        # if flaws is None:\n",
    "        #     continue\n",
    "        sample = fo.Sample(filepath=img)\n",
    "        \n",
    "        # if len(flaws) == 1:\n",
    "        #     continue\n",
    "        for flaw in flaws:\n",
    "            flaw.bbox = flaw.bbox / (img_w, img_h) \n",
    "            detections.append(\n",
    "                fo.Detection(label=flaw.category, bounding_box=flaw.bbox.xywh, confidence=flaw.score)\n",
    "            )\n",
    "        sample[\"predictions\"] = fo.Detections(detections=detections)\n",
    "        samples.append(sample)\n",
    "print(len(samples))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 39,
   "id": "c5089275-e844-44f2-b6cb-831ba08403d0",
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      " 100% |███████████████████| 20/20 [72.7ms elapsed, 0s remaining, 275.1 samples/s]     \n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "['680b585cc8580079777ce7b2',\n",
       " '680b585cc8580079777ce7b3',\n",
       " '680b585cc8580079777ce7b4',\n",
       " '680b585cc8580079777ce7b5',\n",
       " '680b585cc8580079777ce7b6',\n",
       " '680b585cc8580079777ce7b7',\n",
       " '680b585cc8580079777ce7b8',\n",
       " '680b585cc8580079777ce7b9',\n",
       " '680b585cc8580079777ce7ba',\n",
       " '680b585cc8580079777ce7bb',\n",
       " '680b585cc8580079777ce7bc',\n",
       " '680b585cc8580079777ce7bd',\n",
       " '680b585cc8580079777ce7be',\n",
       " '680b585cc8580079777ce7bf',\n",
       " '680b585cc8580079777ce7c0',\n",
       " '680b585cc8580079777ce7c1',\n",
       " '680b585cc8580079777ce7c2',\n",
       " '680b585cc8580079777ce7c3',\n",
       " '680b585cc8580079777ce7c4',\n",
       " '680b585cc8580079777ce7c5']"
      ]
     },
     "execution_count": 39,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "dataset.add_samples(samples)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 40,
   "id": "b88d9410-cd9a-4b97-bce5-3a24ea3e9e91",
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "data": {
      "text/html": [
       "\n",
       "        <iframe\n",
       "            width=\"100%\"\n",
       "            height=\"800\"\n",
       "            src=\"http://0.0.0.0:5151/?notebook=True&subscription=a73a44ba-d0f8-458b-9068-22e824a8fa4e\"\n",
       "            frameborder=\"0\"\n",
       "            allowfullscreen\n",
       "            \n",
       "        ></iframe>\n",
       "        "
      ],
      "text/plain": [
       "<IPython.lib.display.IFrame at 0x7cab265ce4b0>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "session = fo.launch_app(dataset, address=\"0.0.0.0\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "abea62ce-734e-4084-8fc1-e7ce189e7870",
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|███████████████████████████████████████████████████████████████████████████████| 1000/1000 [04:41<00:00,  3.56it/s]\n"
     ]
    }
   ],
   "source": [
    "# export\n",
    "\n",
    "img_save = \"/home/user/ntfs/NIKEDataReview/0424/images\"\n",
    "label_save = \"/home/user/ntfs/NIKEDataReview/0424/labels\"\n",
    "\n",
    "classes = {\"溢胶\": \"0\",\n",
    "           \"线头\":\"1\",\n",
    "           \"脏污\":\"2\"}\n",
    "\n",
    "with open(\"/home/user/ntfs/NIKEDataReview/0424/classes.txt\", \"w\") as f:\n",
    "    for cls in classes.keys():\n",
    "        f.write(cls + \"\\n\")\n",
    "os.makedirs(img_save, exist_ok=True)\n",
    "os.makedirs(label_save, exist_ok=True)\n",
    "for it in tqdm(dataset.view()):\n",
    "    img = it['filepath']\n",
    "    dst = osp.join(img_save, osp.basename(img))\n",
    "    shutil.copyfile(img, dst) \n",
    "    pred_file = img.replace(\".png\", \".json\")\n",
    "    flaws = parse_result(pred_file)\n",
    "    txt = []\n",
    "    if len(flaws):\n",
    "        for flaw in flaws[1:]:\n",
    "            flaw.polygon = flaw.polygon / (img_w, img_h) \n",
    "            flaw.bbox = flaw.bbox / (img_w, img_h) \n",
    "            # print(len(flaw.polygon.flat()))\n",
    "            if len(flaw.polygon.flat()) != 0:\n",
    "                txt.append(classes[flaw.category] + \" \" + \" \".join(map(str, flaw.polygon.flat())))\n",
    "            else:\n",
    "                x1, y1, x2, y2 = flaw.bbox.xyxy\n",
    "                mask = [x1, y1, x1,y2, x2,y2, x2, y1]\n",
    "                txt.append(classes[flaw.category] + \" \" + \" \".join(map(str, mask)))\n",
    "    with open(osp.join(label_save, osp.basename(img).replace('.png', '.txt')), \"w\") as f:\n",
    "        for annotation in txt:\n",
    "            f.write(annotation + \"\\n\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0fad42f3-86ab-4584-aac4-f91ba71823f9",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
